#!/usr/bin/python3
# -*- coding: UTF-8 -*-
# -*- author: Wd0g  -*-

import redis

# 进行redis相关操作
class RedisCache:

    def __init__(self, host='127.0.0.1', port=6379, db=1):
        self.cache = redis.Redis(host= host, port= port, db= db)
        self.domainHashName = "domain"
        self.domainRuleHashName = "domain_rule"
        self.domainSuffixHashName = "domain_suffix"
        self.domainSonInfoHashName = "domain_son_num"
        self.urlMd5HashName = "url_md5"
        self.urlWaitSsetName = "url_wait"
        self.urlLevelHashName = "url_level"

    # 增加域名的爬取数量,并返回增加后的数量,如果为第一次,则返回1
    def addDomainNum(self, domain, num):
        return self.cache.hincrby(self.domainHashName, domain, num)

    # 获取指定域名的扫描数量,如果没有返回0
    def getDomainNum(self, domain):
        num = self.cache.hget(self.domainHashName, domain)
        if num is not None:
            num = int(num)
        else:
            num = 0
        return num

    # 增加需要的域名后缀, 如果是新添加,返回1,否则返回0
    def addDomainSuffix(self, domain):
        return self.cache.hset(self.domainSuffixHashName, domain,1)

    # 返回全部的域名后缀信息
    def getAllDomainSuffix(self):
        domainList = self.cache.hkeys(self.domainSuffixHashName)
        domainList = list(domain.decode(encoding='utf-8') for domain in domainList)
        return domainList

    # 将url的MD5添加到已扫描的HASH表中,如果是新添加,返回1,否则增加后的值
    def addMd5ToHash(self, md5):
        return self.cache.hincrby(self.urlMd5HashName, md5, 1)

    # 从hash表中获取指定key的值,如果有值则返回,否则返回None
    def getMd5ForHash(self, md5):
        value = self.cache.hget(self.urlMd5HashName, md5)
        if value is not None:
            value = int(value)
        return value

    # 添加新的域名规则到表中,如果是新添加,返回1,否则返回0
    def addDomainRule(self, domain, value):
        return self.cache.hset(self.domainRuleHashName, domain, value)

    # 获取当前域名的规则,没有该域名返回0,并将该域名添加到域名规则中
    def getDomainRule(self, domain):
        rule = self.cache.hget(self.domainRuleHashName, domain)
        if rule is not None:
            rule = int(rule)
        else:
            self.addDomainRule(domain, 0)
            rule = 0
        return rule

    # 从域名规则表中删除已经添加到域名规则,成功删除返回1,否则防返回0
    def delDomainRule(self, domain):
        return self.cache.hdel(self.domainRuleHashName, domain)

    # 将url添加到序列中
    def addUrlWait(self, url, level):
        return self.cache.zadd(self.urlWaitSsetName, url,level)

    # 从等待序列中把URL获取出来
    def getUrlWait(self):
        url = self.cache.zrange(self.urlWaitSsetName, 0, 0)  # 获取序列第一个
        if url != []:
            url = url[0].decode(encoding='utf-8')
            self.cache.zremrangebyrank(self.urlWaitSsetName, 0, 0)  # 删除序列第一个
        else:
            url = None
        return url

    # 将URL等级和数量添加到优先级表中,如果是新添加,返回1,否则返回0
    def addUrlLevel(self, level, num):
        return self.cache.hset(self.urlLevelHashName, level, num)

    # 获取所有优先级信息
    def getAllUrlLevel(self):
        infoList = self.cache.hgetall(self.urlLevelHashName)
        newList = {}
        for info in infoList:
            level = int(info)
            num = int(infoList.get(info))
            newList[level] = num
        return newList

    # 增加根域名的子域名数量,并返回增加后的数量
    def addDomainSonNum(self, domain, num):
        return self.cache.hincrby(self.domainSonInfoHashName, domain, num)

    # 获取当前根域名下子域名的数量
    def getDomainSonNum(self, domain):
        num = self.cache.hget(self.domainSonInfoHashName, domain)
        if num is not None:
            num = int(num)
        else:
            num = 0
        return num


class RedisTools:
    def __init__(self):
        self.cache = RedisCache()
        self.urlLevel = self.cache.getAllUrlLevel()                         # URL等级信息
        self.domainSuffixList = self.cache.getAllDomainSuffix()             # 域名后缀信息

        self.domainSonMaxNum = 50  # 子域名最多数量

    # 判断当前URL是否已经扫描过
    def urlIsScan(self, url):
        from function import md5
        urlMd5 = md5(url)
        info = self.cache.getMd5ForHash(urlMd5)
        if info is not None:
            isScan = True
        else:
            isScan = False
        return isScan

    # 将域名添加到域名表中,且增加数量
    def addDomain(self,url):
        from function import md5
        from function import getUrlInfo
        baseDomain, domain, _ = getUrlInfo(url)
        self.cache.addMd5ToHash(md5(url))                   # 将URL的MD5添加到hash表中
        addNum = self.cache.addDomainNum(domain,1)          # 添加域名下的URL扫描数量
        if addNum == 1:
            self.cache.addDomainSonNum(baseDomain, 1)

    # 判断当前域名是否可以进行扫描
    # 1: 检测子域名数量是否超过数量
    # 2: 检测当前域名规则是否可以扫描
    def checkDomain(self, domain):
        # 检测当前域名是否为需要的后缀,如非需要,则直接返回False
        from function import getDomainSuffix
        domainSuffix  = getDomainSuffix(domain, self.domainSuffixList)
        if domainSuffix == None:
            return False


        # 检测当前域名下的子域名数量是否超过设定
        sonNum = self.cache.getDomainSonNum(domain)
        if sonNum > self.domainSonMaxNum:
            self.cache.addDomainRule(domain, 1)         # 设置域名规则为不可扫
            return False

        # 检测当前域名是否可以正常扫描
        rule = self.cache.getDomainRule(domain)
        if rule == 0:
            isOk = True
        else:
            isOk = False
        return isOk

    # 获取当前域名下URL的Level
    def getUrlLevel(self, domain):
        domainScanNum = self.cache.getDomainNum(domain)
        urlLevel = 0
        for index in self.urlLevel:
            if domainScanNum >= self.urlLevel[index]:
                urlLevel = index
        return urlLevel

    # 检测URL数组,写到队列
    def checkUrlList(self, urlList):
        addNum = 0
        if len(urlList) > 0:
            from function import getUrlInfo
            from function import md5

        for url in urlList:
            # 取出根域名,判断根域名是否可以扫描
            baseDomain, domain, _ = getUrlInfo(url)
            if self.checkDomain(baseDomain) is False:
                continue

            # 判断URL是否可以扫描
            urlMd5 = md5(url)
            if self.cache.getMd5ForHash(urlMd5) != None:
                continue
            else:
                urlLevel = self.getUrlLevel(domain)         # 获取域名下URL的level
                self.cache.addUrlWait(url, urlLevel)        # 将URL和URL 优先级写入到等待队列中
            addNum += 1
        return addNum

RedisTools().checkDomain('http://www.gov.com')
